#define K_MAX_SHAPE_DIM 0
#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;
class KernelCopysign {
public:
    __aicore__ inline KernelCopysign() {}
    __aicore__ inline void Init(GM_ADDR input, GM_ADDR other, GM_ADDR out,
                                uint32_t bigNum, uint32_t bigLength, uint32_t smallLength, TPipe* pipeIn) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        uint32_t coreDataNum;
        if (GetBlockIdx() < bigNum) {
            coreDataNum = bigLength;
            inputGm.SetGlobalBuffer((__gm__ DTYPE_INPUT*)input + bigLength * GetBlockIdx(), bigLength);
            otherGm.SetGlobalBuffer((__gm__ DTYPE_INPUT*)other + bigLength * GetBlockIdx(), bigLength);
            outGm.SetGlobalBuffer((__gm__ DTYPE_INPUT*)out + bigLength * GetBlockIdx(), bigLength);
        } else {
            coreDataNum = smallLength;
            inputGm.SetGlobalBuffer((__gm__ DTYPE_INPUT*)input + bigLength * bigNum + smallLength * (GetBlockIdx() - bigNum), smallLength);
            otherGm.SetGlobalBuffer((__gm__ DTYPE_INPUT*)other + bigLength * bigNum + smallLength * (GetBlockIdx() - bigNum), smallLength);
            outGm.SetGlobalBuffer((__gm__ DTYPE_INPUT*)out + bigLength * bigNum + smallLength * (GetBlockIdx() - bigNum), smallLength);
        }
        this->tileDataLenght = 32768 / sizeof(DTYPE_INPUT);
        this->loopCount = coreDataNum/this->tileDataLenght;
        this->tailDataLenght = coreDataNum - this->tileDataLenght*this->loopCount;
        if(this->tailDataLenght == 0){this->tailDataLenght = this->tileDataLenght; this->loopCount--;}
        pipeIn->InitBuffer(inQueueInput, 1, 65792);
        pipeIn->InitBuffer(inQueueOther, 1, 65280);
        pipeIn->InitBuffer(outQueueOut, 1, this->tileDataLenght * sizeof(DTYPE_INPUT));
    }
    __aicore__ inline void Process() {
        this->processDataLenght = this->tileDataLenght;
        for (int32_t i = 0; i < this->loopCount; i++) {
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
        this->processDataLenght = this->tailDataLenght;
        CopyIn(loopCount);
        Compute(loopCount);
        CopyOut(loopCount);
    }
private:
    __aicore__ inline void CopyIn(int32_t progress) {
        LocalTensor<DTYPE_INPUT> inputLocal = inQueueInput.AllocTensor<DTYPE_INPUT>();
        LocalTensor<DTYPE_INPUT> otherLocal = inQueueOther.AllocTensor<DTYPE_INPUT>();
        DataCopy(inputLocal, inputGm[progress * this->tileDataLenght], this->processDataLenght);
        DataCopy(otherLocal, otherGm[progress * this->tileDataLenght], this->processDataLenght);
        inQueueInput.EnQue(inputLocal);
        inQueueOther.EnQue(otherLocal);
    }
    __aicore__ inline void Compute(int32_t progress) {
        LocalTensor<DTYPE_INPUT> inputLocal = inQueueInput.DeQue<DTYPE_INPUT>();
        LocalTensor<DTYPE_INPUT> otherLocal = inQueueOther.DeQue<DTYPE_INPUT>();
        LocalTensor<DTYPE_INPUT> outLocal = outQueueOut.AllocTensor<DTYPE_INPUT>();
        if constexpr (std::is_same_v<DTYPE_INPUT, float>) {
            Abs(inputLocal, inputLocal, this->processDataLenght);
            ShiftRight(otherLocal.ReinterpretCast<uint32_t>(), otherLocal.ReinterpretCast<uint32_t>(), (uint32_t)31, this->processDataLenght);
            ShiftLeft(otherLocal.ReinterpretCast<uint32_t>(), otherLocal.ReinterpretCast<uint32_t>(), (uint32_t)31, this->processDataLenght);
            Or(outLocal.ReinterpretCast<uint16_t>(), inputLocal.ReinterpretCast<uint16_t>(), otherLocal.ReinterpretCast<uint16_t>(), this->processDataLenght*2);
        }
        else
        {
            Abs(inputLocal.ReinterpretCast<half>(), inputLocal.ReinterpretCast<half>(), this->processDataLenght);
            ShiftRight(otherLocal.ReinterpretCast<uint16_t>(), otherLocal.ReinterpretCast<uint16_t>(), (uint16_t)15, this->processDataLenght);
            ShiftLeft(otherLocal.ReinterpretCast<uint16_t>(), otherLocal.ReinterpretCast<uint16_t>(), (uint16_t)15, this->processDataLenght);
            Or(outLocal.ReinterpretCast<uint16_t>(), inputLocal.ReinterpretCast<uint16_t>(), otherLocal.ReinterpretCast<uint16_t>(), this->processDataLenght);
        }
        inQueueInput.FreeTensor(inputLocal);
        inQueueOther.FreeTensor(otherLocal);
        outQueueOut.EnQue<DTYPE_INPUT>(outLocal);
    }
    __aicore__ inline void CopyOut(int32_t progress) {
        LocalTensor<DTYPE_INPUT> outLocal = outQueueOut.DeQue<DTYPE_INPUT>();
        DataCopy(outGm[progress * this->tileDataLenght], outLocal, this->processDataLenght);
        outQueueOut.FreeTensor(outLocal);
    }
private:
    TQue<QuePosition::VECIN, 1> inQueueInput, inQueueOther;
    TQue<QuePosition::VECOUT, 1> outQueueOut;
    TBuf<QuePosition::VECCALC> QueueTmp1, QueueTmp2,QueueTmp3,QueueTmp4, QueueTmp5;
    GlobalTensor<DTYPE_INPUT> inputGm;
    GlobalTensor<DTYPE_INPUT> otherGm;
    GlobalTensor<DTYPE_INPUT> outGm;
    uint32_t loopCount;
    uint32_t tileDataLenght;
    uint32_t tailDataLenght;
    uint32_t processDataLenght;
};
class KernelCopysignBC_1 {
public:
    __aicore__ inline KernelCopysignBC_1() {}
    __aicore__ inline void Init(GM_ADDR input, GM_ADDR other, GM_ADDR out,
                                uint32_t bigNum, uint32_t bigLength, uint32_t smallLength, 
                                uint32_t repeat, uint32_t cycles, uint32_t interval, TPipe* pipeIn) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        uint32_t coreDataNum;
        uint32_t group = GetBlockIdx()/repeat;
        uint32_t within = GetBlockIdx()%repeat;
        if (within < bigNum) {
            coreDataNum = bigLength;
            inputGm.SetGlobalBuffer((__gm__ DTYPE_INPUT*)input + group * cycles * interval + within * bigLength * interval);
            otherGm.SetGlobalBuffer((__gm__ DTYPE_INPUT*)other + group * interval);
            outGm.SetGlobalBuffer((__gm__ DTYPE_INPUT*)out + group * cycles * interval + within * bigLength * interval);
        } else {
            coreDataNum = smallLength;
            inputGm.SetGlobalBuffer((__gm__ DTYPE_INPUT*)input + group * cycles * interval + bigLength * interval * bigNum + smallLength * interval  * (within - bigNum));
            otherGm.SetGlobalBuffer((__gm__ DTYPE_INPUT*)other + group * interval);
            outGm.SetGlobalBuffer((__gm__ DTYPE_INPUT*)out + group * cycles * interval + bigLength * interval * bigNum + smallLength * interval  * (within - bigNum));
        }
        this->cycles = cycles;
        this->interval = interval;
        this->extendNum = 16;
        this->tileDataLenght = this->extendNum * interval;
        this->loopCount = coreDataNum/this->extendNum;
        this->tailDataLenght = (coreDataNum - this->extendNum*this->loopCount) * interval;
        if(this->tailDataLenght == 0){this->tailDataLenght = this->tileDataLenght; this->loopCount--;}
        pipeIn->InitBuffer(inQueueInput, 2, 65792/2);
        pipeIn->InitBuffer(inQueueOther, 2, 65280/2);
        pipeIn->InitBuffer(outQueueOut, BUFFER_NUM, this->tileDataLenght * sizeof(DTYPE_INPUT));
    }
    __aicore__ inline void Process() {
        this->processDataLenght = this->tileDataLenght;
        LocalTensor<DTYPE_INPUT> otherLocal = inQueueOther.AllocTensor<DTYPE_INPUT>();
        DataCopy(otherLocal, otherGm, this->interval);
        inQueueOther.EnQue(otherLocal);
        CopyIn(0);
        otherLocal = inQueueOther.DeQue<DTYPE_INPUT>();
        if constexpr (std::is_same_v<DTYPE_INPUT, float>) {
            ShiftRight(otherLocal.ReinterpretCast<uint32_t>(), otherLocal.ReinterpretCast<uint32_t>(), (uint32_t)31, this->interval);
            ShiftLeft(otherLocal.ReinterpretCast<uint32_t>(), otherLocal.ReinterpretCast<uint32_t>(), (uint32_t)31, this->interval);
            Adds(otherLocal.ReinterpretCast<int32_t>()[this->interval], otherLocal.ReinterpretCast<int32_t>(), (int32_t)0, this->interval);
            Adds(otherLocal.ReinterpretCast<int32_t>()[2*this->interval], otherLocal.ReinterpretCast<int32_t>(), (int32_t)0, 2*this->interval);
            Adds(otherLocal.ReinterpretCast<int32_t>()[4*this->interval], otherLocal.ReinterpretCast<int32_t>(), (int32_t)0, 4*this->interval);
            Adds(otherLocal.ReinterpretCast<int32_t>()[8*this->interval], otherLocal.ReinterpretCast<int32_t>(), (int32_t)0, 8*this->interval);
        }
        else
        {
            ShiftRight(otherLocal.ReinterpretCast<uint16_t>(), otherLocal.ReinterpretCast<uint16_t>(), (uint16_t)15, this->interval);
            ShiftLeft(otherLocal.ReinterpretCast<uint16_t>(), otherLocal.ReinterpretCast<uint16_t>(), (uint16_t)15, this->interval);
            Adds(otherLocal.ReinterpretCast<int16_t>()[this->interval], otherLocal.ReinterpretCast<int16_t>(), (int16_t)0, this->interval);
            Adds(otherLocal.ReinterpretCast<int16_t>()[2*this->interval], otherLocal.ReinterpretCast<int16_t>(), (int16_t)0, 2*this->interval);
            Adds(otherLocal.ReinterpretCast<int16_t>()[4*this->interval], otherLocal.ReinterpretCast<int16_t>(), (int16_t)0, 4*this->interval);
            Adds(otherLocal.ReinterpretCast<int16_t>()[8*this->interval], otherLocal.ReinterpretCast<int16_t>(), (int16_t)0, 8*this->interval);
        }
        Compute(0, otherLocal);
        CopyOut(0);
        for (int32_t i = 1; i < this->loopCount; i++) {
            CopyIn(i);
            Compute(i, otherLocal);
            CopyOut(i);
        }
        this->processDataLenght = this->tailDataLenght;
        CopyIn(loopCount);
        Compute(loopCount, otherLocal);
        CopyOut(loopCount);
        inQueueOther.FreeTensor(otherLocal);
    }
private:
    __aicore__ inline void CopyIn(int32_t progress) {
        LocalTensor<DTYPE_INPUT> inputLocal = inQueueInput.AllocTensor<DTYPE_INPUT>();
        DataCopy(inputLocal, inputGm[progress * this->tileDataLenght], this->processDataLenght);
        inQueueInput.EnQue(inputLocal);
    }
    __aicore__ inline void Compute(int32_t progress, LocalTensor<DTYPE_INPUT> otherLocal) {
        LocalTensor<DTYPE_INPUT> inputLocal = inQueueInput.DeQue<DTYPE_INPUT>();
        LocalTensor<DTYPE_INPUT> outLocal = outQueueOut.AllocTensor<DTYPE_INPUT>();
        if constexpr (std::is_same_v<DTYPE_INPUT, float>) {
            Abs(inputLocal, inputLocal, this->processDataLenght);
            Or(outLocal.ReinterpretCast<uint16_t>(), inputLocal.ReinterpretCast<uint16_t>(), otherLocal.ReinterpretCast<uint16_t>(), this->processDataLenght*2);
        }
        else
        {
            Abs(inputLocal.ReinterpretCast<half>(), inputLocal.ReinterpretCast<half>(), this->processDataLenght);
            Or(outLocal.ReinterpretCast<uint16_t>(), inputLocal.ReinterpretCast<uint16_t>(), otherLocal.ReinterpretCast<uint16_t>(), this->processDataLenght);
        }
        inQueueInput.FreeTensor(inputLocal);
        outQueueOut.EnQue<DTYPE_INPUT>(outLocal);
    }
    __aicore__ inline void CopyOut(int32_t progress) {
        LocalTensor<DTYPE_INPUT> outLocal = outQueueOut.DeQue<DTYPE_INPUT>();
        DataCopy(outGm[progress * this->tileDataLenght], outLocal, this->processDataLenght);
        outQueueOut.FreeTensor(outLocal);
    }
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueInput, inQueueOther;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueOut;
    TBuf<QuePosition::VECCALC> QueueTmp1, QueueTmp2,QueueTmp3,QueueTmp4, QueueTmp5;
    GlobalTensor<DTYPE_INPUT> inputGm;
    GlobalTensor<DTYPE_INPUT> otherGm;
    GlobalTensor<DTYPE_INPUT> outGm;
    uint32_t cycles;
    uint32_t interval;
    uint32_t loopCount;
    uint32_t tileDataLenght;
    uint32_t tailDataLenght;
    uint32_t processDataLenght;
    uint32_t extendNum;
};
extern "C" __global__ __aicore__ void copysign(GM_ADDR input, GM_ADDR other, GM_ADDR out, GM_ADDR workspace, GM_ADDR tiling) {
    TPipe pipe;
    if(TILING_KEY_IS(1)) { 
        GET_TILING_DATA_WITH_STRUCT(CopysignTilingData, tiling_data, tiling);
        KernelCopysign op;
        op.Init(input, other, out,
                tiling_data.bigNum, tiling_data.bigLength, tiling_data.smallLength, &pipe);
        op.Process();
    }
    else if(TILING_KEY_IS(3)) {
        GET_TILING_DATA_WITH_STRUCT(CopysignTilingDataBC, tiling_data, tiling);
        KernelCopysignBC_1 op;
        op.Init(input, other, out,
                tiling_data.bigNum, tiling_data.bigLength, tiling_data.smallLength, 
                tiling_data.repeat, tiling_data.cycles, tiling_data.interval, &pipe);
        op.Process();
    }
}